"""
====================================================================
MongoDB集合初始化工具 / MongoDB Collection Initialization
====================================================================
功能说明：
1. 创建新集合和索引
2. 验证集合是否存在
3. 提供回滚和清理功能

使用方法：
    python -c "from utils.mongo_init import init_all_collections; init_all_collections()"

作者: Data Analysis Team
日期: 2025-11-06
====================================================================
"""

import sys
from typing import List, Tuple
from utils.mongo_util import db
from model.db_schemas import COLLECTIONS_CONFIG

# 颜色输出（用于终端显示）
class Colors:
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    BLUE = '\033[94m'
    END = '\033[0m'


def create_index(collection_name: str, index_spec: List[Tuple]) -> bool:
    """
    创建索引

    参数:
        collection_name: 集合名称
        index_spec: 索引规范，例如 [("memberId", 1), ("period_start", 1)]

    返回:
        bool: 是否成功
    """
    try:
        collection = db[collection_name]
        index_name = collection.create_index(index_spec)
        print(f"  {Colors.GREEN}✓{Colors.END} 创建索引: {index_spec} -> {index_name}")
        return True
    except Exception as e:
        print(f"  {Colors.RED}✗{Colors.END} 索引创建失败: {e}")
        return False


def create_ttl_index(collection_name: str, field_name: str, expire_after_seconds: int = 0) -> bool:
    """
    创建TTL索引（自动过期删除）

    参数:
        collection_name: 集合名称
        field_name: 时间戳字段名
        expire_after_seconds: 过期时间（秒），0表示文档ttl字段值到期即删除

    返回:
        bool: 是否成功
    """
    try:
        collection = db[collection_name]
        collection.create_index([(field_name, 1)], expireAfterSeconds=expire_after_seconds)
        print(f"  {Colors.GREEN}✓{Colors.END} 创建TTL索引: {field_name} (expireAfterSeconds={expire_after_seconds})")
        return True
    except Exception as e:
        print(f"  {Colors.RED}✗{Colors.END} TTL索引创建失败: {e}")
        return False


def collection_exists(collection_name: str) -> bool:
    """检查集合是否存在"""
    return collection_name in db.list_collection_names()


def get_collection_info(collection_name: str) -> dict:
    """获取集合信息"""
    if not collection_exists(collection_name):
        return {}

    collection = db[collection_name]
    return {
        "name": collection_name,
        "documents": collection.count_documents({}),
        "size": db.command("collstats", collection_name).get("size", 0),
        "indices": list(collection.list_indexes())
    }


def init_collection(collection_name: str, force: bool = False) -> bool:
    """
    初始化单个集合

    参数:
        collection_name: 集合名称
        force: 是否强制重建（删除现有集合）

    返回:
        bool: 是否成功
    """
    print(f"\n{Colors.BLUE}初始化集合: {collection_name}{Colors.END}")

    config = COLLECTIONS_CONFIG.get(collection_name)
    if not config:
        print(f"  {Colors.RED}✗{Colors.END} 集合配置不存在")
        return False

    # 检查集合是否已存在
    if collection_exists(collection_name):
        print(f"  ℹ 集合已存在，现有文档数: {db[collection_name].count_documents({})}")
        if not force:
            print(f"  ℹ 跳过创建（使用 force=True 可强制重建）")
            return True

    # 创建集合
    try:
        db.create_collection(collection_name)
        print(f"  {Colors.GREEN}✓{Colors.END} 集合创建成功")
    except Exception as e:
        if "already exists" not in str(e):
            print(f"  {Colors.RED}✗{Colors.END} 集合创建失败: {e}")
            return False

    # 创建索引
    if config['indices']:
        print(f"  创建索引 ({len(config['indices'])} 个):")

        # 特殊处理TTL索引
        for idx_spec in config['indices']:
            if isinstance(idx_spec, dict):
                # TTL索引
                field_name = idx_spec.get('field')
                expire_seconds = idx_spec.get('expireAfterSeconds', 0)
                create_ttl_index(collection_name, field_name, expire_seconds)
            elif isinstance(idx_spec, tuple) and len(idx_spec) == 2:
                # 普通索引或TTL索引
                if isinstance(idx_spec[1], dict) and 'expireAfterSeconds' in idx_spec[1]:
                    # TTL索引：[("ttl", 1), {"expireAfterSeconds": 0}]
                    field_name = idx_spec[0][0] if isinstance(idx_spec[0], tuple) else idx_spec[0]
                    expire_seconds = idx_spec[1].get('expireAfterSeconds', 0)
                    create_ttl_index(collection_name, field_name, expire_seconds)
                else:
                    create_index(collection_name, idx_spec)
            else:
                create_index(collection_name, idx_spec)
    else:
        print(f"  ℹ 无需创建索引")

    return True


def init_all_collections(force: bool = False) -> bool:
    """
    初始化所有集合

    参数:
        force: 是否强制重建现有集合

    返回:
        bool: 全部初始化是否成功
    """
    print(f"\n{Colors.BLUE}{'='*70}")
    print("MongoDB集合初始化")
    print(f"{'='*70}{Colors.END}")

    success_count = 0
    total_count = len(COLLECTIONS_CONFIG)

    for collection_name in COLLECTIONS_CONFIG.keys():
        if init_collection(collection_name, force):
            success_count += 1

    print(f"\n{Colors.BLUE}{'='*70}")
    print(f"初始化完成: {success_count}/{total_count} 集合成功")
    print(f"{'='*70}{Colors.END}\n")

    return success_count == total_count


def cleanup_collection(collection_name: str, confirm: bool = True) -> bool:
    """
    清理集合（删除所有文档）

    参数:
        collection_name: 集合名称
        confirm: 是否要求确认

    返回:
        bool: 是否成功
    """
    if not collection_exists(collection_name):
        print(f"{Colors.YELLOW}⚠{Colors.END} 集合不存在: {collection_name}")
        return False

    collection = db[collection_name]
    doc_count = collection.count_documents({})

    if confirm:
        response = input(f"\n确定要删除 {collection_name} 中的 {doc_count} 个文档吗? (yes/no): ")
        if response.lower() != "yes":
            print("已取消")
            return False

    try:
        result = collection.delete_many({})
        print(f"{Colors.GREEN}✓{Colors.END} 已删除 {result.deleted_count} 个文档")
        return True
    except Exception as e:
        print(f"{Colors.RED}✗{Colors.END} 删除失败: {e}")
        return False


def drop_collection(collection_name: str, confirm: bool = True) -> bool:
    """
    删除整个集合

    参数:
        collection_name: 集合名称
        confirm: 是否要求确认

    返回:
        bool: 是否成功
    """
    if not collection_exists(collection_name):
        print(f"{Colors.YELLOW}⚠{Colors.END} 集合不存在: {collection_name}")
        return False

    if confirm:
        response = input(f"\n确定要删除集合 {collection_name} 吗? (yes/no): ")
        if response.lower() != "yes":
            print("已取消")
            return False

    try:
        db.drop_collection(collection_name)
        print(f"{Colors.GREEN}✓{Colors.END} 集合已删除: {collection_name}")
        return True
    except Exception as e:
        print(f"{Colors.RED}✗{Colors.END} 删除失败: {e}")
        return False


def show_collection_stats():
    """显示所有集合的统计信息"""
    print(f"\n{Colors.BLUE}MongoDB集合统计信息{Colors.END}")
    print(f"{'='*70}")

    for collection_name in COLLECTIONS_CONFIG.keys():
        info = get_collection_info(collection_name)
        if info:
            size_mb = info['size'] / (1024 * 1024)
            print(f"\n集合: {Colors.BLUE}{collection_name}{Colors.END}")
            print(f"  文档数: {info['documents']}")
            print(f"  大小: {size_mb:.2f} MB")
            print(f"  索引数: {len(info['indices']) - 1}")  # 减1因为默认有_id索引
        else:
            print(f"\n集合: {Colors.RED}{collection_name}{Colors.END} (不存在或无权限)")

    print(f"{'='*70}\n")


# ====================================================================
# 命令行工具
# ====================================================================

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="MongoDB集合初始化工具")
    parser.add_argument("--init", action="store_true", help="初始化所有集合")
    parser.add_argument("--force", action="store_true", help="强制重建现有集合")
    parser.add_argument("--stats", action="store_true", help="显示集合统计信息")
    parser.add_argument("--cleanup", type=str, help="清理指定集合中的数据")
    parser.add_argument("--drop", type=str, help="删除指定集合")
    parser.add_argument("--skip-confirm", action="store_true", help="跳过确认提示")

    args = parser.parse_args()

    if args.init:
        init_all_collections(force=args.force)
    elif args.stats:
        show_collection_stats()
    elif args.cleanup:
        cleanup_collection(args.cleanup, confirm=not args.skip_confirm)
    elif args.drop:
        drop_collection(args.drop, confirm=not args.skip_confirm)
    else:
        parser.print_help()
